"""
06_indexing_and_constraints.py
==============================
Neo4j Indexing and Constraints

This module covers:
1. Range indexes for fast lookups
2. Composite indexes
3. Text indexes for search
4. Fulltext indexes
5. Vector indexes (HNSW)
6. Unique constraints
7. Existence constraints
8. Node key constraints
9. Index performance analysis

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum
import time
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class IndexType(Enum):
    """Types of indexes available in Neo4j."""
    RANGE = "RANGE"           # B-tree for equality/range
    TEXT = "TEXT"             # For string operations
    POINT = "POINT"           # For spatial data
    FULLTEXT = "FULLTEXT"     # Lucene-based search
    VECTOR = "VECTOR"         # HNSW for embeddings


class ConstraintType(Enum):
    """Types of constraints available."""
    UNIQUE = "UNIQUE"
    EXISTS = "EXISTS"           # Enterprise only
    NODE_KEY = "NODE_KEY"       # Enterprise only


@dataclass
class IndexInfo:
    """Information about an index."""
    name: str
    type: str
    labels: List[str]
    properties: List[str]
    state: str
    population_percent: float


class IndexManager:
    """
    Manages indexes and constraints in Neo4j.
    
    Index Types and Use Cases:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                    INDEX SELECTION GUIDE                        │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  RANGE INDEX                                                    │
    │  ───────────                                                    │
    │  Use for: = , <> , < , > , <= , >= , IN, STARTS WITH           │
    │  Best for: Equality lookups, range queries                      │
    │  Example: WHERE n.id = 'P001'                                  │
    │                                                                 │
    │  TEXT INDEX                                                     │
    │  ──────────                                                     │
    │  Use for: CONTAINS, ENDS WITH                                   │
    │  Best for: String searches within values                        │
    │  Example: WHERE n.name CONTAINS 'Smith'                        │
    │                                                                 │
    │  FULLTEXT INDEX                                                 │
    │  ──────────────                                                 │
    │  Use for: Full-text search with scoring                         │
    │  Best for: Natural language search                              │
    │  Example: CALL db.index.fulltext.queryNodes()                  │
    │                                                                 │
    │  VECTOR INDEX (HNSW)                                           │
    │  ───────────────────                                            │
    │  Use for: Similarity search on embeddings                       │
    │  Best for: Semantic search, RAG applications                   │
    │  Example: CALL db.index.vector.queryNodes()                    │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, driver=None):
        self.driver = driver or get_driver()
    
    # =========================================================
    # RANGE INDEXES
    # =========================================================
    
    def create_range_index(
        self,
        label: str,
        property_name: str,
        index_name: str = None
    ) -> str:
        """
        Create a range index for fast property lookups.
        
        Range indexes use B-tree structure:
        - O(log n) lookup time
        - Supports equality and range queries
        - Memory efficient
        
        Args:
            label: Node label to index
            property_name: Property to index
            index_name: Optional custom index name
            
        Returns:
            Index name
        """
        name = index_name or f"idx_{label.lower()}_{property_name.lower()}"
        
        query = f"""
        CREATE INDEX {name} IF NOT EXISTS
        FOR (n:{label})
        ON (n.{property_name})
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    def create_composite_index(
        self,
        label: str,
        properties: List[str],
        index_name: str = None
    ) -> str:
        """
        Create a composite index on multiple properties.
        
        Useful when queries filter on multiple properties together.
        Order matters: put most selective property first.
        
        Args:
            label: Node label
            properties: List of properties (order matters)
            index_name: Optional index name
            
        Returns:
            Index name
        """
        name = index_name or f"idx_{label.lower()}_composite"
        props_str = ", ".join([f"n.{p}" for p in properties])
        
        query = f"""
        CREATE INDEX {name} IF NOT EXISTS
        FOR (n:{label})
        ON ({props_str})
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    # =========================================================
    # TEXT INDEXES
    # =========================================================
    
    def create_text_index(
        self,
        label: str,
        property_name: str,
        index_name: str = None
    ) -> str:
        """
        Create a text index for string searches.
        
        Optimized for:
        - CONTAINS
        - ENDS WITH
        
        Note: STARTS WITH uses range index, not text index.
        
        Args:
            label: Node label
            property_name: String property
            index_name: Optional index name
            
        Returns:
            Index name
        """
        name = index_name or f"idx_{label.lower()}_{property_name.lower()}_text"
        
        query = f"""
        CREATE TEXT INDEX {name} IF NOT EXISTS
        FOR (n:{label})
        ON (n.{property_name})
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    # =========================================================
    # FULLTEXT INDEXES
    # =========================================================
    
    def create_fulltext_index(
        self,
        index_name: str,
        labels: List[str],
        properties: List[str],
        analyzer: str = "standard"
    ) -> str:
        """
        Create a fulltext index for natural language search.
        
        Features:
        - Tokenization
        - Stemming
        - Fuzzy matching
        - Relevance scoring
        
        Analyzers:
        - standard: Default, good for most languages
        - simple: Lowercase only
        - whitespace: Split on whitespace
        - keyword: No tokenization
        
        Args:
            index_name: Index name
            labels: Node labels to index
            properties: Properties to include
            analyzer: Text analyzer
            
        Returns:
            Index name
        """
        labels_str = "|".join(labels)
        props_str = ", ".join([f"n.{p}" for p in properties])
        
        query = f"""
        CREATE FULLTEXT INDEX {index_name} IF NOT EXISTS
        FOR (n:{labels_str})
        ON EACH [{props_str}]
        OPTIONS {{
            indexConfig: {{
                `fulltext.analyzer`: '{analyzer}'
            }}
        }}
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return index_name
    
    def fulltext_search(
        self,
        index_name: str,
        search_string: str,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Search using fulltext index.
        
        Search syntax supports:
        - Single terms: diabetes
        - Phrases: "type 2 diabetes"
        - Wildcards: diabet*
        - Fuzzy: diabetes~
        - Boolean: diabetes AND treatment
        """
        query = """
        CALL db.index.fulltext.queryNodes($index_name, $search)
        YIELD node, score
        RETURN 
            labels(node)[0] as label,
            node.id as id,
            node.name as name,
            score
        ORDER BY score DESC
        LIMIT $limit
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                search=search_string,
                limit=limit
            )
            return [r.data() for r in result]
    
    # =========================================================
    # VECTOR INDEXES
    # =========================================================
    
    def create_vector_index(
        self,
        label: str,
        property_name: str,
        dimensions: int,
        similarity_function: str = "cosine",
        index_name: str = None,
        m: int = 16,
        ef_construction: int = 64
    ) -> str:
        """
        Create an HNSW vector index for similarity search.
        
        HNSW Parameters:
        - dimensions: Vector size (must match embeddings)
        - similarity_function: 'cosine' or 'euclidean'
        - m: Max connections per node (higher = better recall, more memory)
        - ef_construction: Build-time beam width (higher = better quality)
        
        Args:
            label: Node label
            property_name: Property containing vectors
            dimensions: Vector dimensions
            similarity_function: Distance metric
            index_name: Optional index name
            m: HNSW M parameter
            ef_construction: HNSW construction parameter
            
        Returns:
            Index name
        """
        name = index_name or f"idx_{label.lower()}_vector"
        
        query = f"""
        CREATE VECTOR INDEX {name} IF NOT EXISTS
        FOR (n:{label})
        ON (n.{property_name})
        OPTIONS {{
            indexConfig: {{
                `vector.dimensions`: {dimensions},
                `vector.similarity_function`: '{similarity_function}'
            }}
        }}
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    def vector_search(
        self,
        index_name: str,
        query_vector: List[float],
        top_k: int = 10,
        min_score: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Perform vector similarity search.
        
        Args:
            index_name: Vector index name
            query_vector: Query embedding
            top_k: Number of results
            min_score: Minimum similarity (0-1 for cosine)
            
        Returns:
            Similar nodes with scores
        """
        query = """
        CALL db.index.vector.queryNodes($index_name, $k, $vector)
        YIELD node, score
        WHERE score >= $min_score
        RETURN 
            node.id as id,
            node.content as content,
            score
        ORDER BY score DESC
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_vector,
                min_score=min_score
            )
            return [r.data() for r in result]
    
    # =========================================================
    # CONSTRAINTS
    # =========================================================
    
    def create_unique_constraint(
        self,
        label: str,
        property_name: str,
        constraint_name: str = None
    ) -> str:
        """
        Create a unique constraint on a property.
        
        This:
        1. Ensures uniqueness of values
        2. Automatically creates an index
        3. Prevents duplicate creation
        
        Args:
            label: Node label
            property_name: Property to make unique
            constraint_name: Optional constraint name
            
        Returns:
            Constraint name
        """
        name = constraint_name or f"constraint_{label.lower()}_{property_name.lower()}_unique"
        
        query = f"""
        CREATE CONSTRAINT {name} IF NOT EXISTS
        FOR (n:{label})
        REQUIRE n.{property_name} IS UNIQUE
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    def create_existence_constraint(
        self,
        label: str,
        property_name: str,
        constraint_name: str = None
    ) -> str:
        """
        Create an existence constraint (Enterprise only).
        
        Ensures the property must be set on all nodes with this label.
        
        Args:
            label: Node label
            property_name: Required property
            constraint_name: Optional constraint name
            
        Returns:
            Constraint name
        """
        name = constraint_name or f"constraint_{label.lower()}_{property_name.lower()}_exists"
        
        query = f"""
        CREATE CONSTRAINT {name} IF NOT EXISTS
        FOR (n:{label})
        REQUIRE n.{property_name} IS NOT NULL
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    def create_node_key(
        self,
        label: str,
        properties: List[str],
        constraint_name: str = None
    ) -> str:
        """
        Create a node key constraint (Enterprise only).
        
        Combines:
        - Uniqueness on combination of properties
        - Existence for all properties
        
        Args:
            label: Node label
            properties: Properties that form the key
            constraint_name: Optional constraint name
            
        Returns:
            Constraint name
        """
        name = constraint_name or f"constraint_{label.lower()}_key"
        props_str = ", ".join([f"n.{p}" for p in properties])
        
        query = f"""
        CREATE CONSTRAINT {name} IF NOT EXISTS
        FOR (n:{label})
        REQUIRE ({props_str}) IS NODE KEY
        """
        
        with self.driver.session() as session:
            session.run(query)
        
        return name
    
    # =========================================================
    # INDEX MANAGEMENT
    # =========================================================
    
    def list_indexes(self) -> List[IndexInfo]:
        """List all indexes in the database."""
        query = """
        SHOW INDEXES
        YIELD name, type, labelsOrTypes, properties, state, populationPercent
        RETURN name, type, labelsOrTypes as labels, properties, state, 
               populationPercent
        ORDER BY name
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [
                IndexInfo(
                    name=r["name"],
                    type=r["type"],
                    labels=r["labels"] or [],
                    properties=r["properties"] or [],
                    state=r["state"],
                    population_percent=r["populationPercent"] or 0
                )
                for r in result
            ]
    
    def list_constraints(self) -> List[Dict[str, Any]]:
        """List all constraints in the database."""
        query = """
        SHOW CONSTRAINTS
        YIELD name, type, labelsOrTypes, properties, ownedIndex
        RETURN name, type, labelsOrTypes as labels, properties, ownedIndex
        ORDER BY name
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    def drop_index(self, index_name: str) -> bool:
        """Drop an index by name."""
        query = f"DROP INDEX {index_name} IF EXISTS"
        
        with self.driver.session() as session:
            session.run(query)
            return True
    
    def drop_constraint(self, constraint_name: str) -> bool:
        """Drop a constraint by name."""
        query = f"DROP CONSTRAINT {constraint_name} IF EXISTS"
        
        with self.driver.session() as session:
            session.run(query)
            return True
    
    def wait_for_index(
        self,
        index_name: str,
        timeout_seconds: int = 60
    ) -> bool:
        """
        Wait for index to become ONLINE.
        
        Args:
            index_name: Index to wait for
            timeout_seconds: Maximum wait time
            
        Returns:
            True if online, False if timeout
        """
        query = """
        SHOW INDEXES
        YIELD name, state
        WHERE name = $name
        RETURN state
        """
        
        start_time = time.time()
        
        with self.driver.session() as session:
            while time.time() - start_time < timeout_seconds:
                result = session.run(query, name=index_name)
                record = result.single()
                
                if record and record["state"] == "ONLINE":
                    return True
                
                time.sleep(0.5)
        
        return False
    
    # =========================================================
    # PERFORMANCE ANALYSIS
    # =========================================================
    
    def analyze_query(self, query: str, params: Dict = None) -> Dict[str, Any]:
        """
        Analyze query execution plan with PROFILE.
        
        PROFILE actually executes the query and shows:
        - Rows processed at each step
        - DB hits (I/O operations)
        - Time spent
        - Index usage
        """
        profile_query = f"PROFILE {query}"
        
        with self.driver.session() as session:
            result = session.run(profile_query, params or {})
            summary = result.consume()
            
            return {
                "db_hits": summary.counters.nodes_created,  # Simplified
                "plan": str(summary.plan) if summary.plan else "No plan available",
                "available_after": summary.result_available_after,
                "consumed_after": summary.result_consumed_after
            }
    
    def explain_query(self, query: str) -> str:
        """
        Get query execution plan without executing.
        
        EXPLAIN shows what Neo4j would do:
        - Planned operators
        - Estimated rows
        - Index usage
        """
        explain_query = f"EXPLAIN {query}"
        
        with self.driver.session() as session:
            result = session.run(explain_query)
            summary = result.consume()
            
            if summary.plan:
                return self._format_plan(summary.plan)
            return "No plan available"
    
    def _format_plan(self, plan, indent=0) -> str:
        """Format execution plan for display."""
        lines = []
        prefix = "  " * indent
        
        operator = plan.operator_type
        args = plan.arguments
        estimated_rows = args.get("EstimatedRows", "?")
        
        lines.append(f"{prefix}├─ {operator} (est. {estimated_rows} rows)")
        
        if plan.children:
            for child in plan.children:
                lines.append(self._format_plan(child, indent + 1))
        
        return "\n".join(lines)


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Indexing and Constraints - Demonstration")
    print("=" * 60)
    
    try:
        manager = IndexManager()
        
        # Create indexes
        print("\n📌 Creating Indexes")
        print("-" * 40)
        
        # Range index
        idx1 = manager.create_range_index("Patient", "id")
        print(f"   Created range index: {idx1}")
        
        # Composite index
        idx2 = manager.create_composite_index("Patient", ["gender", "active"])
        print(f"   Created composite index: {idx2}")
        
        # Text index
        idx3 = manager.create_text_index("Patient", "name")
        print(f"   Created text index: {idx3}")
        
        # Fulltext index
        idx4 = manager.create_fulltext_index(
            "patient_search",
            ["Patient", "Disease"],
            ["name", "description"]
        )
        print(f"   Created fulltext index: {idx4}")
        
        # Constraints
        print("\n📌 Creating Constraints")
        print("-" * 40)
        
        c1 = manager.create_unique_constraint("Patient", "id")
        print(f"   Created unique constraint: {c1}")
        
        # List indexes
        print("\n📌 Listing Indexes")
        print("-" * 40)
        
        indexes = manager.list_indexes()
        for idx in indexes[:5]:  # Show first 5
            print(f"   {idx.name}: {idx.type} on {idx.labels} - {idx.state}")
        
        print("\n" + "=" * 60)
        print("✅ Indexing demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
